using System.Collections;
using System.Collections.Generic;
using System.Runtime.InteropServices;
using UnityEngine;
using UnityEngine.Rendering;

namespace GPUParticles
{
public class ParticlesManager : MonoBehaviour
{
    private bool debug=true;
    public string VERSION = "1001";
    
    #region Module Variables
    #region General
    //派遣线程
    public const int THREAD_COUNT = 256;
    public int maxParticles = 256;
    public float timeScale = 1f;
    #endregion
    #region Color
    
    public Color color = Color.white;
    public Gradient colorOverLife;
    public int colorSteps = 16;
    private Texture2D colorOverLifeTexture;
    #endregion

    #region Size
    
    public float size = 1f;
    // public AnimationCurve sizeOverLife;
    // private ComputeBuffer sizeOverLifeBuffer;
    // public int sizeSteps = 16;
    #endregion
    
    #region Lifetime
    public float minLifetime = 3f;
    public float maxLifetime = 5f;
    #endregion
    #region Inherit Info
    public bool enableInheritVelocity = true;
    public float inheritVelocity = 0.3f;
    public float extrapolation = 0f;
    private Vector3 previousPositon;
    #endregion
    #region Emission
    public bool enableEmission = true;
    public float emissionRate = 1500f;
    public float minInitialSpeed = 0.1f;
    public float maxInitialSpeed = 0.5f;
    
    public float sphereRadius = 0.2f;
    
    public Vector3 direction;
    #endregion
    #region Assets
    public ComputeShader computeShader;
    public Material renderMaterial;
    public Material lineRenderMaterial;
    #endregion
    #region Effect
    [Range(0,5)]
    public float max_linelen = 2;

    private ComputeBuffer line_index,line_Counter;
    private int effect_kernel;
    private int []effect_debug_array;
    private int line_buffersize;
    #endregion
    #endregion
    
    
    #region Other Variables

    private int initKernel, emitKernel, updateKernel,
        dispatchkernel,rendingdispatchkernel;
    private struct Particle
    {
        public bool alive;
        public Vector3 position;
        public Vector3 velocity;
        public Vector2 life; //x = age, y = lifetime
        public Color color;
        float size;
    }
    private ComputeBuffer particles,
        particle_pool,
        quad, 
        counter,
        UpdateCounter,
        Simulation_Counter,
        Rendingargs,
        debug_buffer
        
        ;

    
    private ComputeBuffer[] alive_buffer;
    private int bufferSize;
    private int groupCount;
    private int[] counterArray;
    private int[] debug_Array;
    private int deadCount = 0;
    private int aliveCount = 0;
    private uint pre_alive_index=0;
    private uint now_alive_index=1;
    private uint Render_index=1;
    #endregion
    
    #region Unity Functions

    private void Awake()
    {
        Camera.main.depthTextureMode = DepthTextureMode.Depth;

        DispatchInit();
    }

    private void Update()
    {
        
        
        DispatchEmit(Mathf.RoundToInt(Time.deltaTime * emissionRate * timeScale));
        
        DispatchUpdate();
    }

    private void OnRenderObject()
    {
        DispatchRending();
        {
           
            ComputeBuffer.CopyCount(alive_buffer[Render_index],Simulation_Counter,0);
        
            computeShader.SetBuffer(dispatchkernel,"UpdateIndirectArgumentBuffer",UpdateCounter);
            computeShader.SetBuffer(dispatchkernel,"Simulation_Counter",Simulation_Counter);
        
            computeShader.Dispatch(dispatchkernel,1,1,1);
        }
        {
            computeShader.SetBuffer(effect_kernel, "particles", particles);
            
            computeShader.SetBuffer(effect_kernel, "Particle_alive_Counter", Simulation_Counter);
            computeShader.SetBuffer(effect_kernel, "alive_index", alive_buffer[Render_index]);
            
            computeShader.SetBuffer(effect_kernel, "line_index", line_index);
            computeShader.SetBuffer(effect_kernel, "line_Counter", line_Counter);
            
            computeShader.SetFloat("line_len",max_linelen);
            
            computeShader.DispatchIndirect(effect_kernel,UpdateCounter);
            // if (debug)
            // {
            //     line_index.GetData(debug_Array);
            // }
        }
        
        renderMaterial.SetBuffer("particles", particles);
        renderMaterial.SetBuffer("nowalive", alive_buffer[Render_index]);
     
        renderMaterial.SetBuffer("quad", quad);
        renderMaterial.SetPass(0);
        int count = GetAliveCount();
        SetDeadCount();
        if (count != (bufferSize - GetDeadCount()))
        {
            Debug.Log("Wrong!"+count+"dead"+deadCount);
        }
        
        //Graphics.DrawProceduralNow(MeshTopology.Triangles, 6, count);
        Graphics.DrawProceduralIndirectNow(MeshTopology.Triangles,Rendingargs,0);
        
        
        
        computeShader.SetBuffer(rendingdispatchkernel,"Rending_Counter",line_Counter);
        computeShader.SetBuffer(rendingdispatchkernel,"DrawIndirectArgumentBuffer",Rendingargs);
        computeShader.SetInt("Render_flag",1);
        computeShader.Dispatch(rendingdispatchkernel,1,1,1);
        
        lineRenderMaterial.SetBuffer("particles", particles);
        lineRenderMaterial.SetBuffer("nowalive", line_index);
        
        lineRenderMaterial.SetPass(0);
        Graphics.DrawProceduralIndirectNow(MeshTopology.Lines,Rendingargs,0);
        
        
        
        line_index.SetCounterValue(0);
        line_Counter.SetCounterValue(0);
        
    }

    private void OnDestroy()
    {
        ReleaseBuffers();
    }

    #endregion
   
    #region Public Functions

    public int GetAliveCount()
    {
        return aliveCount;
    }

    public void EffectInit()
    {
        effect_kernel=computeShader.FindKernel("CalLine");
        line_buffersize = bufferSize * 6;
        line_index = new ComputeBuffer(line_buffersize, sizeof(uint), ComputeBufferType.Append);
        line_Counter = new ComputeBuffer(1, sizeof(uint), ComputeBufferType.Counter);
        line_Counter.SetCounterValue(0);
        line_index.SetCounterValue(0);
        effect_debug_array = new int[line_buffersize];
    }

    public void EffectDestroy()
    {
        if(line_index!=null)line_index.Release();
    }
    public void DispatchInit()
    {
        ReleaseBuffers();

      
        

        initKernel = computeShader.FindKernel("Init");
        emitKernel = computeShader.FindKernel("Emit");
        updateKernel = computeShader.FindKernel("Update");
        dispatchkernel = computeShader.FindKernel("Dispatch");
        rendingdispatchkernel = computeShader.FindKernel("DispatchDraw");
        groupCount = Mathf.CeilToInt((float)maxParticles / THREAD_COUNT);
        bufferSize = groupCount * THREAD_COUNT;

        particles = new ComputeBuffer(bufferSize, Marshal.SizeOf(typeof(Particle)));
        Simulation_Counter = new ComputeBuffer(1, sizeof(uint), ComputeBufferType.Raw);
        debug_buffer = new ComputeBuffer(256, sizeof(uint), ComputeBufferType.Raw);
        debug_Array = new int[256];
        particle_pool = new ComputeBuffer(bufferSize, sizeof(int), ComputeBufferType.Append);
        particle_pool.SetCounterValue(0);
        alive_buffer = new ComputeBuffer[2]
        {
            new ComputeBuffer(bufferSize, sizeof(int), ComputeBufferType.Append),
            new ComputeBuffer(bufferSize, sizeof(int), ComputeBufferType.Append)
        };
        alive_buffer[pre_alive_index].SetCounterValue(0);
        alive_buffer[now_alive_index].SetCounterValue(0);
        counter = new ComputeBuffer(4, sizeof(int), ComputeBufferType.IndirectArguments);
        
        Rendingargs = new ComputeBuffer(1, sizeof(uint)*4, ComputeBufferType.IndirectArguments);
        
        counterArray = new int[] { 0, 1, 0, 0 };
        UpdateCounter = new ComputeBuffer(1, sizeof(uint)*3, ComputeBufferType.IndirectArguments);

        EffectInit();
        computeShader.SetBuffer(initKernel, "particles", particles);
        computeShader.SetBuffer(initKernel, "append", particle_pool);

        computeShader.Dispatch(initKernel, groupCount, 1, 1);
       
        
        
        
        quad = new ComputeBuffer(6, Marshal.SizeOf(typeof(Vector3)));
        quad.SetData(new[]
        {
            // new Vector3(-0.5f,-0.5f),
            // new Vector3(-0.5f,0.5f),
            // new Vector3(0.5f,0.5f),
            // new Vector3(0.5f,-0.5f),
            new Vector3(-0.5f,0.5f),
            new Vector3(0.5f,0.5f),
            new Vector3(0.5f,-0.5f),
            new Vector3(0.5f,-0.5f),
            new Vector3(-0.5f,-0.5f),
            new Vector3(-0.5f,0.5f)
        });
    }
    public void DispatchEmit(int count)
    {
        if (enableEmission)
        {
            count = Mathf.Min(count, maxParticles - (aliveCount));
            

            if (count > 0)
            {
          
                Vector3 velocity = (transform.position - previousPositon) / Time.deltaTime;
                previousPositon = transform.position;
               
                computeShader.SetBuffer(emitKernel, "particles", particles);
               
                computeShader.SetBuffer(emitKernel, "Particle_Pool", particle_pool);
               
                computeShader.SetBuffer(emitKernel, "Particle_alive", alive_buffer[pre_alive_index]);
                
             
                

                computeShader.SetVector("seeds",
                    new Vector3(Random.Range(1f, 10000f), Random.Range(1f, 10000f), Random.Range(1f, 10000f)));
                computeShader.SetVector("initialSpeedRange", new Vector2(minInitialSpeed, maxInitialSpeed));
                computeShader.SetVector("inheritedPosition", transform.position);
                
                computeShader.SetVector("lifeRange", new Vector2(minLifetime, maxLifetime));
                computeShader.SetVector("time", new Vector2(Time.deltaTime, Time.time));
                computeShader.SetVector("color", color);
                computeShader.SetFloat("size", size);
                
                
                
                computeShader.SetFloat("radius", Mathf.Max(0.01f, sphereRadius));
                
                if (enableInheritVelocity)
                {
                    computeShader.SetVector("inheritedVelocity", velocity * inheritVelocity);
                    computeShader.SetFloat("extrapolation", extrapolation);
                }
                computeShader.Dispatch(emitKernel, count, 1, 1);
                
                //computeShader.DispatchIndirect(emitKernel,UpdateCounter);
            }
        }
        //
        ComputeBuffer.CopyCount(alive_buffer[pre_alive_index],Simulation_Counter,0);
       
        computeShader.SetBuffer(dispatchkernel,"UpdateIndirectArgumentBuffer",UpdateCounter);
        computeShader.SetBuffer(dispatchkernel,"Simulation_Counter",Simulation_Counter);
        
        computeShader.Dispatch(dispatchkernel,1,1,1);
        
        
        
        
        
    }
    #endregion
  

    #region Private Functions

        private int GetDeadCount()
        {
            return deadCount;
        }

        private void SetAliveCount()
        {
            //counter.SetData(counterArray);
            ComputeBuffer.CopyCount(alive_buffer[Render_index], counter, 0);
            counter.GetData(counterArray);
            aliveCount= counterArray[0];
        }
        private void SetDeadCount()
        {
            if (alive_buffer[pre_alive_index] == null 
                ||alive_buffer[now_alive_index]==null 
                ||counter == null || counterArray == null
                ||particle_pool==null)
            {
                deadCount = bufferSize;
                return;
            }
            ComputeBuffer.CopyCount(particle_pool, counter, 0);
            counter.GetData(counterArray);
            deadCount = counterArray[0];
        }

        private void IndirectArgumentBuffers()
        {
            
        }

        private void DispatchRending()
        {
            ComputeBuffer.CopyCount(alive_buffer[Render_index],Simulation_Counter,0);
            computeShader.SetBuffer(rendingdispatchkernel,"Rending_Counter",Simulation_Counter);
            computeShader.SetBuffer(rendingdispatchkernel,"DrawIndirectArgumentBuffer",Rendingargs);
            computeShader.SetInt("Render_flag",0);
            computeShader.Dispatch(rendingdispatchkernel,1,1,1);
            
        }
        private void DispatchUpdate()
        {
            
            
            //alive_buffer[pre_alive_index].GetData(debug_Array);
            if (timeScale > 0)
            {
                //IndirectArgumentBuffers();
                computeShader.SetBuffer(updateKernel, "particles", particles);
                computeShader.SetBuffer(updateKernel, "test", debug_buffer);
                computeShader.SetBuffer(updateKernel,"Simulation_Counter",Simulation_Counter);
                computeShader.SetBuffer(updateKernel, "UpdateIndirectArgumentBuffer", UpdateCounter);
                computeShader.SetBuffer(updateKernel, "Particle_pool_Update", particle_pool);
                computeShader.SetBuffer(updateKernel, "Particle_alive_Update", alive_buffer[pre_alive_index]);
                computeShader.SetBuffer(updateKernel, "Particle_alive_Update_Append", alive_buffer[now_alive_index]);
                
                //computeShader.SetBuffer(updateKernel, "Particle_alive_Update_Append", nowbuf);
                computeShader.SetFloat("size", size);
                
                computeShader.SetVector("time", new Vector2(Time.deltaTime * timeScale, Time.time));
                
                computeShader.DispatchIndirect(updateKernel,UpdateCounter);
                //debug_buffer.GetData(debug_Array);
                Render_index = now_alive_index;
                SetAliveCount();
                if (now_alive_index == 0)
                {
                    now_alive_index = 1;
                    pre_alive_index = 0;
                }
                else
                {
                    now_alive_index = 0;
                    pre_alive_index = 1;
                }
               
                
            }
            
        }

        private void ReleaseBuffers()
        {
            if (particles != null) particles.Release();
            if (alive_buffer!=null)
            {
                if (alive_buffer[pre_alive_index] != null) alive_buffer[pre_alive_index].Release();
                if (alive_buffer[now_alive_index] != null) alive_buffer[now_alive_index].Release();
            }
            if(Rendingargs!=null)Rendingargs.Release();
            if(debug_buffer!=null)debug_buffer.Release();
            if(Simulation_Counter!=null)Simulation_Counter.Release();
            if (particle_pool != null) particle_pool.Release();
            if (counter != null) counter.Release();
            if (quad != null) quad.Release();
            if(UpdateCounter!=null)UpdateCounter.Release();
            EffectDestroy();
        }

        #endregion

   

   

    
}
}
